//
//  visualizer.cpp
//  yuzon
//
//  Created by Alvin on 2023/3/29.
//

#include "visualizer.hpp"

using namespace cv;
using namespace std;

void VisualizeResults(Mat &image,
                      Mat &annotation_image,
                      const HumanResults &human_results,
                      const KeypointResults &keypoint_results,
                      bool draw_bbox,
                      bool draw_pose,
                      int bbox_line_thickness,
                      int pose_keypoint_num,
                      int pose_circle_radius,
                      int pose_line_thickness) {
  cvtColor(image, image, COLOR_RGB2RGBA);
  annotation_image = Mat::zeros(image.size(), CV_8UC4);
  
  if (draw_bbox && !human_results.empty())
    DrawBbox(image, annotation_image, human_results, bbox_line_thickness);
  
  if (draw_pose && !keypoint_results.empty())
    DrawPose(image, annotation_image, keypoint_results,
             pose_keypoint_num, pose_circle_radius, pose_line_thickness);
}

void DrawBbox(Mat &image,
              Mat &annotation_image,
              const HumanResults &results,
              int line_thickness) {
  for (const HumanResult &result : results) {
    string text = result.class_name + " " + to_string(int(result.score * 100)) + "%";
    int font_face = FONT_HERSHEY_PLAIN;
    double font_scale = 1.5;
    float font_thickness = 1.0f;
    Size text_size = getTextSize(text, font_face, font_scale, font_thickness, nullptr);
    
    rectangle(image, result.rect, result.color, line_thickness);
    rectangle(annotation_image, result.rect, result.color, line_thickness);
    rectangle(image,
              Point2d(result.rect.x, result.rect.y - round(text_size.height * 1.25f)),
              Point2d(result.rect.x + result.rect.width, result.rect.y),
              result.color, -1);
    rectangle(annotation_image,
              Point2d(result.rect.x, result.rect.y - round(text_size.height * 1.25f)),
              Point2d(result.rect.x + result.rect.width, result.rect.y),
              result.color, -1);
    putText(image, text, Point2d(result.rect.x, result.rect.y),
            font_face, font_scale, Scalar(255, 255, 255, 255), font_thickness);
    putText(annotation_image, text, Point2d(result.rect.x, result.rect.y),
            font_face, font_scale, Scalar(255, 255, 255, 255), font_thickness);
  }
}

void DrawPose(Mat &image,
              Mat &annotation_image,
              const KeypointResults &results,
              int keypoint_num,
              int circle_radius,
              int line_thickness) {
  const KeypointResults &keypoint_results = (keypoint_num == 17 ?
                                             results :
                                             ConvertKeypoints(results));
  const vector<vector<int>> &edges = kEdges[keypoint_num];
  const vector<int> &color_indexes = kColorIndexes[keypoint_num];
  
  for (const KeypointResult &result : keypoint_results) {
    for (int i = 0; i < edges.size(); ++i) {
      if (result.keypoints[edges[i][0] * 3] < kKeypointThreshold ||
          result.keypoints[edges[i][1] * 3] < kKeypointThreshold)
        continue;
      
      int x1 = round(result.keypoints[edges[i][0] * 3 + 1]);
      int y1 = round(result.keypoints[edges[i][0] * 3 + 2]);
      int x2 = round(result.keypoints[edges[i][1] * 3 + 1]);
      int y2 = round(result.keypoints[edges[i][1] * 3 + 2]);
      line(image, Point2d(x1, y1), Point2d(x2, y2),
           kColorMap[color_indexes[i]], line_thickness);
      line(annotation_image, Point2d(x1, y1), Point2d(x2, y2),
           kColorMap[color_indexes[i]], line_thickness);
    }
    
    for (int i = 0; i < result.num_joints; ++i) {
      if (result.keypoints[i * 3] < kKeypointThreshold)
        continue;
      
      int x = round(result.keypoints[i * 3 + 1]);
      int y = round(result.keypoints[i * 3 + 2]);
      circle(image, Point2d(x, y), circle_radius, kColorMap[8], -1);
      circle(annotation_image, Point2d(x, y), circle_radius, kColorMap[8], -1);
    }
  }
}

KeypointResults ConvertKeypoints(const KeypointResults &results) {
  KeypointResults results_14;
  
  for (const KeypointResult &result : results) {
    KeypointResult result_14;
    result_14.num_joints = 14;
    
    for (int i = 0; i < result.num_joints; ++i) {
      if (i == 0 || i > 4) {
        result_14.keypoints.emplace_back(result.keypoints[i * 3]);
        result_14.keypoints.emplace_back(result.keypoints[i * 3 + 1]);
        result_14.keypoints.emplace_back(result.keypoints[i * 3 + 2]);
      }
      
      if (i == 1) {
        result_14.keypoints.emplace_back((result.keypoints[5 * 3] +
                                          result.keypoints[6 * 3]) / 2.0f);
        result_14.keypoints.emplace_back((result.keypoints[5 * 3 + 1] +
                                          result.keypoints[6 * 3 + 1]) / 2.0f);
        result_14.keypoints.emplace_back((result.keypoints[5 * 3 + 2] +
                                          result.keypoints[6 * 3 + 2]) / 2.0f);
      }
    }
    
    results_14.emplace_back(result_14);
  }
  
  return results_14;
}
